package avltree;

/**
 * Created with IntelliJ IDEA.
 * Description:
 * User: a
 * Date: 2022-11-08
 * Time: 23:09
 */
public class AVLTree {
    static class TreeNode{
        public int val;
        public int bf;  //平衡因子
        public TreeNode left;
        public TreeNode right;
        public TreeNode parent;

        public TreeNode(int val){
            this.val = val;
        }
    }

    public TreeNode root;  //根节点

    public boolean insert(int val){
        TreeNode node = new TreeNode(val);
        if(root == null){
            root = node;
            return true;
        }
        //接下来的步骤类似于二叉搜索树的插入
        TreeNode parent = null;
        TreeNode cur = root;
        while(cur != null){
            if(cur.val < val){
                parent = cur;
                cur = cur.right;
            }else if(cur.val == val){
                return false;
            }else{
                parent = cur;
                cur = cur.left;
            }
        }
        if(parent.val < val){
            parent.right = node;
        }else{
            parent.left = node;
        }
        node.parent = parent;
        cur = node;
        //平衡因子的修改(调节)
        while(parent != null){
            //先看cur是parent的左孩子还是右孩子, 决定平衡因子是++还是--(此处规定的计算平衡因子规则: 右树高度-左树高度)
            if(cur == parent.right){
                parent.bf++;
            }else{
                parent.bf--;
            }
            if(parent.bf == 0){
                //说明已经平衡了
                break;
            }else if(parent.bf == 1 || parent.bf == -1){
                //继续向上进行调整平衡因子
                cur = parent;
                parent = cur.parent;
            }else{
                if(parent.bf == 2){
                    if(cur.bf == 1){
                        //左旋
                        rotateLeft(parent);
                    }else{
                        //cur.bf == -1
                        //右左旋
                        rotateRL(parent);
                    }
                }else{
                    //parent.bf == -2
                    if(cur.bf == -1){
                        //右旋
                        rotateRight(parent);
                    }else{
                        //cur.bf == 1
                        //左右旋
                        rotateLR(parent);
                    }
                }
                break;
            }
        }
        return true;
    }

    //左单旋
    private void rotateLeft(TreeNode parent){
        TreeNode subR = parent.right;
        TreeNode subRL = subR.left;
        TreeNode pParent = parent.parent;
        parent.right = subRL;
        //subRL存在的时候才执行
        if(subRL != null){
            subRL.parent = parent;
        }
        subR.left = parent;
        parent.parent = subR;

        if(parent == root){
            root = subR;
            root.parent = null;
        }else{
            if(pParent.left == parent){
                pParent.left = subR;
            }else{
                pParent.right = subR;
            }
            subR.parent = pParent;
        }
        subR.bf = 0;
        parent.bf = 0;
    }

    //右单旋
    private void rotateRight(TreeNode parent){
        TreeNode subL = parent.left;
        TreeNode subLR = subL.right;
        TreeNode pParent = parent.parent;  //记录parent的父节点
        parent.left = subLR;
        //subLR存在的时候才执行
        if(subLR != null) {
            subLR.parent = parent;
        }
        subL.right = parent;
        parent.parent = subL;

        if(parent == root){
            subL = root;
            root.parent = null;
        }else{
            if(pParent.left == parent){
                pParent.left = subL;
            }else{
                pParent.right = subL;
            }
            subL.parent = pParent;
        }
        subL.bf = 0;
        parent.bf = 0;
    }

    //左右双旋
    private void rotateLR(TreeNode parent){
        TreeNode subL = parent.left;
        TreeNode subLR = subL.right;
        int bf = subLR.bf;

        rotateLeft(parent.left);
        rotateRight(parent);

        if(bf == -1){
            subL.bf = 0;
            subLR.bf = 0;
            parent.bf = 1;
        }else if(bf == 1){
            subL.bf = -1;
            subLR.bf = 0;
            parent.bf = 0;
        }
    }

    //右左双旋
    private void rotateRL(TreeNode parent){
        TreeNode subR = parent.right;
        TreeNode subRL = subR.left;
        int bf = subRL.bf;

        rotateRight(parent.right);
        rotateLeft(parent);

        if(bf == 1){
            parent.bf = -1;
            subR.bf = 0;
            subRL.bf = 0;
        }else if(bf == -1){
            parent.bf = 0;
            subR.bf = 1;
            subRL.bf = 0;
        }
    }

    //中序遍历(判断是否是二叉搜索树)
    public void inorder(TreeNode root){
        if(root == null){
            return;
        }
        inorder(root.left);
        System.out.println(root.val + " ");
        inorder(root.right);
    }

    //获取当前节点的左右子树的高度
    private int height(TreeNode root){
        if(root == null){
            return 0;
        }
        int leftH = height(root.left);
        int rightH = height(root.right);
        return leftH > rightH ? leftH + 1 : rightH + 1;
    }

    //判断当前节点的树是否平衡
    public boolean isBalanced(TreeNode root){
        if(root == null){
            return true;
        }
        int leftH = height(root.left);
        int rightH = height(root.right);
        if(rightH - leftH != root.bf){
            System.out.println("节点:" + root.val + "的平衡因子错误");
            return false;
        }
        return Math.abs(leftH - rightH) <= 1 && isBalanced(root.left) && isBalanced(root.right);
    }
}
